package main

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log"
	"sync"
	"time"
)

// 连接池结构体
type ConnectionPool struct {
	connections chan *sql.DB
	factory     func() (*sql.DB, error)
	mu          sync.RWMutex
	closed      bool
	maxSize     int
}

var (
	ErrPoolClosed  = errors.New("连接池已关闭")
	ErrPoolTimeout = errors.New("获取连接超时")
)

// 创建连接池
func NewConnectionPool(factory func() (*sql.DB, error), maxSize int) (*ConnectionPool, error) {
	if maxSize <= 0 {
		return nil, errors.New("连接池大小必须大于0")
	}

	pool := &ConnectionPool{
		connections: make(chan *sql.DB, maxSize),
		factory:     factory,
		maxSize:     maxSize,
		closed:      false,
	}

	// 预先创建连接
	for i := 0; i < maxSize/2; i++ {
		conn, err := factory()
		if err != nil {
			return nil, err
		}
		pool.connections <- conn
	}

	return pool, nil
}

// 获取连接（带超时）
func (p *ConnectionPool) Get(timeout time.Duration) (*sql.DB, error) {
	p.mu.RLock()
	if p.closed {
		p.mu.RUnlock()
		return nil, ErrPoolClosed
	}
	p.mu.RUnlock()

	select {
	case conn := <-p.connections:
		// 检查连接是否有效
		if err := p.ping(conn); err != nil {
			// 连接无效，创建新连接
			newConn, err := p.factory()
			if err != nil {
				return nil, err
			}
			return newConn, nil
		}
		return conn, nil

	case <-time.After(timeout):
		return nil, ErrPoolTimeout

	default:
		// 如果没有可用连接且未达上限，创建新连接
		p.mu.Lock()
		defer p.mu.Unlock()

		if len(p.connections) < p.maxSize {
			conn, err := p.factory()
			if err != nil {
				return nil, err
			}
			return conn, nil
		}

		// 等待连接释放
		select {
		case conn := <-p.connections:
			return conn, nil
		case <-time.After(timeout):
			return nil, ErrPoolTimeout
		}
	}
}

// 释放连接
func (p *ConnectionPool) Put(conn *sql.DB) error {
	p.mu.RLock()
	if p.closed {
		p.mu.RUnlock()
		return conn.Close()
	}
	p.mu.RUnlock()

	// 检查连接是否有效
	if err := p.ping(conn); err != nil {
		return conn.Close() // 无效连接直接关闭
	}

	select {
	case p.connections <- conn:
		return nil
	default:
		// 连接池已满，关闭连接
		return conn.Close()
	}
}

// 检查连接是否有效
func (p *ConnectionPool) ping(conn *sql.DB) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	return conn.PingContext(ctx)
}

// 关闭连接池
func (p *ConnectionPool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.closed {
		return
	}

	p.closed = true
	close(p.connections)

	// 关闭所有连接
	for conn := range p.connections {
		conn.Close()
	}
}

// 获取连接池状态
func (p *ConnectionPool) Status() (int, int) {
	p.mu.RLock()
	defer p.mu.RUnlock()

	if p.closed {
		return 0, 0
	}

	current := len(p.connections)
	return current, p.maxSize
}

func main() {
	// 创建连接工厂函数
	factory := func() (*sql.DB, error) {
		db, err := sql.Open("mysql", "user:password@tcp(localhost:3306)/test")
		if err != nil {
			return nil, err
		}

		// 设置连接参数
		db.SetMaxOpenConns(1)
		db.SetMaxIdleConns(1)
		db.SetConnMaxLifetime(30 * time.Minute)

		return db, nil
	}

	// 创建连接池
	pool, err := NewConnectionPool(factory, 10)
	if err != nil {
		log.Fatal(err)
	}
	defer pool.Close()

	// 并发使用示例
	var wg sync.WaitGroup
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()

			// 获取连接
			conn, err := pool.Get(5 * time.Second)
			if err != nil {
				log.Printf("协程 %d 获取连接失败: %v", id, err)
				return
			}
			defer pool.Put(conn)

			// 使用连接执行查询
			var result string
			err = conn.QueryRow("SELECT 'Hello from goroutine ' || ?", id).Scan(&result)
			if err != nil {
				log.Printf("查询失败: %v", err)
				return
			}

			fmt.Printf("协程 %d: %s\n", id, result)
		}(i)
	}

	wg.Wait()

	// 查看连接池状态
	current, max := pool.Status()
	fmt.Printf("连接池状态: %d/%d 连接\n", current, max)
}
